import asyncio
import aiohttp
import json
import time
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field
from abc import ABC, abstractmethod


@dataclass
class AIResponse:
    """AI响应结果"""
    provider: str
    success: bool
    response: str
    error: Optional[str] = None
    response_time: float = 0.0
    metadata: Dict[str, Any] = field(default_factory=dict)


class BaseAIProvider(ABC):
    """AI服务提供商基类"""
    
    def __init__(self, api_key: str, model: str = None):
        self.api_key = api_key
        self.model = model
        self.name = self.__class__.__name__.replace('Provider', '')
    
    @abstractmethod
    async def send_message(self, message: str, session: aiohttp.ClientSession) -> AIResponse:
        """发送消息到AI服务"""
        pass


class BaiduErnieProvider(BaseAIProvider):
    """百度文心一言服务提供商"""
    
    def __init__(self, api_key: str, secret_key: str, model: str = "ernie-bot"):
        super().__init__(api_key, model)
        self.secret_key = secret_key
        self.access_token = None
        self.token_expires = 0
    
    async def get_access_token(self, session: aiohttp.ClientSession) -> str:
        """获取百度API访问令牌"""
        if self.access_token and time.time() < self.token_expires:
            return self.access_token
            
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {
            'grant_type': 'client_credentials',
            'client_id': self.api_key,
            'client_secret': self.secret_key
        }
        
        async with session.post(url, params=params) as response:
            data = await response.json()
            self.access_token = data['access_token']
            self.token_expires = time.time() + data['expires_in'] - 300  # 提前5分钟过期
            return self.access_token
    
    async def send_message(self, message: str, session: aiohttp.ClientSession) -> AIResponse:
        start_time = time.time()
        try:
            access_token = await self.get_access_token(session)
            url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/{self.model}?access_token={access_token}"
            
            payload = {
                "messages": [
                    {"role": "user", "content": message}
                ]
            }
            
            async with session.post(url, json=payload) as response:
                data = await response.json()
                response_time = time.time() - start_time
                
                if 'result' in data:
                    return AIResponse(
                        provider=self.name,
                        success=True,
                        response=data['result'],
                        response_time=response_time,
                        metadata={"usage": data.get('usage', {})}
                    )
                else:
                    return AIResponse(
                        provider=self.name,
                        success=False,
                        response="",
                        error=data.get('error_msg', '未知错误'),
                        response_time=response_time
                    )
                    
        except Exception as e:
            return AIResponse(
                provider=self.name,
                success=False,
                response="",
                error=str(e),
                response_time=time.time() - start_time
            )


class AliyunQwenProvider(BaseAIProvider):
    """阿里云通义千问服务提供商"""
    
    def __init__(self, api_key: str, model: str = "qwen-plus"):
        super().__init__(api_key, model)
    
    async def send_message(self, message: str, session: aiohttp.ClientSession) -> AIResponse:
        start_time = time.time()
        try:
            url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
            
            headers = {
                'Authorization': f'Bearer {self.api_key}',
                'Content-Type': 'application/json'
            }
            
            payload = {
                "model": self.model,
                "input": {
                    "messages": [
                        {"role": "user", "content": message}
                    ]
                }
            }
            
            async with session.post(url, json=payload, headers=headers) as response:
                data = await response.json()
                response_time = time.time() - start_time
                
                if 'output' in data and 'text' in data['output']:
                    return AIResponse(
                        provider=self.name,
                        success=True,
                        response=data['output']['text'],
                        response_time=response_time,
                        metadata={"usage": data.get('usage', {})}
                    )
                else:
                    return AIResponse(
                        provider=self.name,
                        success=False,
                        response="",
                        error=data.get('message', '未知错误'),
                        response_time=response_time
                    )
                    
        except Exception as e:
            return AIResponse(
                provider=self.name,
                success=False,
                response="",
                error=str(e),
                response_time=time.time() - start_time
            )


class ZhipuGLMProvider(BaseAIProvider):
    """智谱AI GLM服务提供商"""
    
    def __init__(self, api_key: str, model: str = "glm-4"):
        super().__init__(api_key, model)
    
    async def send_message(self, message: str, session: aiohttp.ClientSession) -> AIResponse:
        start_time = time.time()
        try:
            url = "https://open.bigmodel.cn/api/paas/v4/chat/completions"
            
            headers = {
                'Authorization': f'Bearer {self.api_key}',
                'Content-Type': 'application/json'
            }
            
            payload = {
                "model": self.model,
                "messages": [
                    {"role": "user", "content": message}
                ]
            }
            
            async with session.post(url, json=payload, headers=headers) as response:
                data = await response.json()
                response_time = time.time() - start_time
                
                if 'choices' in data and len(data['choices']) > 0:
                    return AIResponse(
                        provider=self.name,
                        success=True,
                        response=data['choices'][0]['message']['content'],
                        response_time=response_time,
                        metadata={"usage": data.get('usage', {})}
                    )
                else:
                    return AIResponse(
                        provider=self.name,
                        success=False,
                        response="",
                        error=data.get('error', {}).get('message', '未知错误'),
                        response_time=response_time
                    )
                    
        except Exception as e:
            return AIResponse(
                provider=self.name,
                success=False,
                response="",
                error=str(e),
                response_time=time.time() - start_time
            )


class KimiProvider(BaseAIProvider):
    """月之暗面Kimi服务提供商"""
    
    def __init__(self, api_key: str, model: str = "moonshot-v1-8k"):
        super().__init__(api_key, model)
    
    async def send_message(self, message: str, session: aiohttp.ClientSession) -> AIResponse:
        start_time = time.time()
        try:
            url = "https://api.moonshot.cn/v1/chat/completions"
            
            headers = {
                'Authorization': f'Bearer {self.api_key}',
                'Content-Type': 'application/json'
            }
            
            payload = {
                "model": self.model,
                "messages": [
                    {"role": "user", "content": message}
                ]
            }
            
            async with session.post(url, json=payload, headers=headers) as response:
                data = await response.json()
                response_time = time.time() - start_time
                
                if 'choices' in data and len(data['choices']) > 0:
                    return AIResponse(
                        provider=self.name,
                        success=True,
                        response=data['choices'][0]['message']['content'],
                        response_time=response_time,
                        metadata={"usage": data.get('usage', {})}
                    )
                else:
                    return AIResponse(
                        provider=self.name,
                        success=False,
                        response="",
                        error=data.get('error', {}).get('message', '未知错误'),
                        response_time=response_time
                    )
                    
        except Exception as e:
            return AIResponse(
                provider=self.name,
                success=False,
                response="",
                error=str(e),
                response_time=time.time() - start_time
            )


class DeepSeekProvider(BaseAIProvider):
    """DeepSeek服务提供商"""
    
    def __init__(self, api_key: str, model: str = "deepseek-chat"):
        super().__init__(api_key, model)
    
    async def send_message(self, message: str, session: aiohttp.ClientSession) -> AIResponse:
        start_time = time.time()
        try:
            url = "https://api.deepseek.com/v1/chat/completions"
            
            headers = {
                'Authorization': f'Bearer {self.api_key}',
                'Content-Type': 'application/json'
            }
            
            payload = {
                "model": self.model,
                "messages": [
                    {"role": "user", "content": message}
                ]
            }
            
            async with session.post(url, json=payload, headers=headers) as response:
                data = await response.json()
                response_time = time.time() - start_time
                
                if 'choices' in data and len(data['choices']) > 0:
                    return AIResponse(
                        provider=self.name,
                        success=True,
                        response=data['choices'][0]['message']['content'],
                        response_time=response_time,
                        metadata={"usage": data.get('usage', {})}
                    )
                else:
                    return AIResponse(
                        provider=self.name,
                        success=False,
                        response="",
                        error=data.get('error', {}).get('message', '未知错误'),
                        response_time=response_time
                    )
                    
        except Exception as e:
            return AIResponse(
                provider=self.name,
                success=False,
                response="",
                error=str(e),
                response_time=time.time() - start_time
            )


class MultiAIChat:
    """多AI对话管理器"""
    
    def __init__(self):
        self.providers: List[BaseAIProvider] = []
    
    def add_provider(self, provider: BaseAIProvider):
        """添加AI服务提供商"""
        self.providers.append(provider)
    
    def remove_provider(self, provider_name: str):
        """移除AI服务提供商"""
        self.providers = [p for p in self.providers if p.name != provider_name]
    
    async def send_to_all(self, message: str, timeout: int = 30) -> List[AIResponse]:
        """向所有AI服务发送消息"""
        if not self.providers:
            return []
        
        timeout_obj = aiohttp.ClientTimeout(total=timeout)
        async with aiohttp.ClientSession(timeout=timeout_obj) as session:
            tasks = [
                provider.send_message(message, session) 
                for provider in self.providers
            ]
            
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 处理异常结果
            responses = []
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    responses.append(AIResponse(
                        provider=self.providers[i].name,
                        success=False,
                        response="",
                        error=str(result),
                        response_time=0.0
                    ))
                else:
                    responses.append(result)
            
            return responses
    
    def print_results(self, responses: List[AIResponse]):
        """格式化打印所有响应结果"""
        print(f"\n{'='*80}")
        print(f"AI响应结果汇总 (共{len(responses)}个服务)")
        print(f"{'='*80}")
        
        successful_responses = [r for r in responses if r.success]
        failed_responses = [r for r in responses if not r.success]
        
        print(f"\n✅ 成功响应: {len(successful_responses)} 个")
        print(f"❌ 失败响应: {len(failed_responses)} 个")
        
        # 显示成功的响应
        for response in successful_responses:
            print(f"\n{'-'*60}")
            print(f"🤖 {response.provider}")
            print(f"⏱️  响应时间: {response.response_time:.2f}秒")
            print(f"📝 回答:\n{response.response}")
            if response.metadata.get('usage'):
                usage = response.metadata['usage']
                print(f"📊 使用情况: {usage}")
        
        # 显示失败的响应
        if failed_responses:
            print(f"\n{'-'*60}")
            print("❌ 失败的服务:")
            for response in failed_responses:
                print(f"   • {response.provider}: {response.error}")
    
    def get_summary(self, responses: List[AIResponse]) -> Dict[str, Any]:
        """获取响应结果摘要"""
        successful = [r for r in responses if r.success]
        failed = [r for r in responses if not r.success]
        
        return {
            "total_providers": len(responses),
            "successful_count": len(successful),
            "failed_count": len(failed),
            "success_rate": len(successful) / len(responses) if responses else 0,
            "average_response_time": sum(r.response_time for r in successful) / len(successful) if successful else 0,
            "fastest_provider": min(successful, key=lambda x: x.response_time).provider if successful else None,
            "slowest_provider": max(successful, key=lambda x: x.response_time).provider if successful else None,
            "successful_providers": [r.provider for r in successful],
            "failed_providers": [r.provider for r in failed]
        }


# 便捷函数
def create_multi_ai_chat() -> MultiAIChat:
    """创建多AI对话管理器实例"""
    return MultiAIChat()


# 示例配置函数
def setup_example_providers(chat: MultiAIChat, config: Dict[str, Dict[str, str]]):
    """根据配置设置AI提供商
    
    config格式示例:
    {
        "baidu": {"api_key": "your_api_key", "secret_key": "your_secret_key"},
        "aliyun": {"api_key": "your_api_key"},
        "zhipu": {"api_key": "your_api_key"},
        "kimi": {"api_key": "your_api_key"},
        "deepseek": {"api_key": "your_api_key"}
    }
    """
    if "baidu" in config:
        provider = BaiduErnieProvider(
            api_key=config["baidu"]["api_key"],
            secret_key=config["baidu"]["secret_key"]
        )
        chat.add_provider(provider)
    
    if "aliyun" in config:
        provider = AliyunQwenProvider(api_key=config["aliyun"]["api_key"])
        chat.add_provider(provider)
    
    if "zhipu" in config:
        provider = ZhipuGLMProvider(api_key=config["zhipu"]["api_key"])
        chat.add_provider(provider)
    
    if "kimi" in config:
        provider = KimiProvider(api_key=config["kimi"]["api_key"])
        chat.add_provider(provider)
    
    if "deepseek" in config:
        provider = DeepSeekProvider(api_key=config["deepseek"]["api_key"])
        chat.add_provider(provider)


if __name__ == "__main__":
    # 示例使用
    async def main():
        # 创建多AI对话管理器
        chat = create_multi_ai_chat()
        
        # 注意：这里需要你自己的API密钥
        config = {
            # "baidu": {"api_key": "your_baidu_api_key", "secret_key": "your_secret_key"},
            # "aliyun": {"api_key": "your_aliyun_api_key"},
            # "zhipu": {"api_key": "your_zhipu_api_key"},
            # "kimi": {"api_key": "your_kimi_api_key"},
            # "deepseek": {"api_key": "your_deepseek_api_key"}
        }
        
        # setup_example_providers(chat, config)
        
        if not chat.providers:
            print("⚠️ 没有配置任何AI服务提供商，请在config中添加你的API密钥")
            return
        
        # 发送测试消息
        question = "你好，请简单介绍一下人工智能的发展历程。"
        print(f"📤 发送问题: {question}")
        
        responses = await chat.send_to_all(question)
        chat.print_results(responses)
        
        # 打印摘要
        summary = chat.get_summary(responses)
        print(f"\n📊 响应摘要:")
        print(f"   成功率: {summary['success_rate']:.1%}")
        print(f"   平均响应时间: {summary['average_response_time']:.2f}秒")
        if summary['fastest_provider']:
            print(f"   最快服务: {summary['fastest_provider']}")
    
    # 运行示例
    # asyncio.run(main())
    print("🎯 多AI对话模块已加载完成！")
    print("💡 使用说明:")
    print("   1. 配置你的API密钥")
    print("   2. 调用 setup_example_providers() 设置服务商")
    print("   3. 使用 chat.send_to_all() 发送消息") 